import matplotlib.pyplot as plt
import numpy as np


def is_valid_center(new_center, centers, min_distance):
    for center in centers:
        distance = np.linalg.norm(new_center - center)
        if distance < min_distance:
            return False
    return True


def generate_main_centers(count):
    main_centers = [np.random.rand(2) * 10]
    for _ in range(count - 1):
        while True:
            new_main_center = np.random.rand(2) * 30
            if is_valid_center(new_main_center, main_centers, 12):
                main_centers.append(new_main_center)
                break
    return main_centers


def generate_word_cloud(keyword_list, tendency_list, class_list=None):
    if class_list is None:
        class_list = np.random.randint(1, 5, len(keyword_list))
    class_dict = {}
    main_center_classes = []
    for class_item in class_list:
        if class_item not in class_dict:
            class_dict[class_item] = 1
            main_center_classes.append(class_item)
        else:
            class_dict[class_item] += 1
    count = len(class_dict)

    plt.figure(figsize=(5, 5), facecolor="#eeeeee")

    np.random.seed(21)
    # centers = [np.random.rand(2) * 10]
    main_centers = generate_main_centers(count)
    final_centers = []
    for i, value in enumerate(class_dict.values()):
        centers = [main_centers[i]]
        # print(value)
        for j in range(value - 1):
            while True:
                new_center = np.random.rand(2) * 10
                if is_valid_center(new_center, centers, 3):
                    centers.append(new_center)
                    break
        # print(centers)
        final_centers.append(centers)

    # for _ in range(len(keyword_list) - 1):
    #     while True:
    #         new_center = np.random.rand(2) * 10
    #         if is_valid_center(new_center, centers, 4):
    #             centers.append(new_center)
    #             break
    #
    radius = [1 + tendency for tendency in tendency_list]
    colors = plt.cm.viridis(np.linspace(0, 1, len(keyword_list)))

    for i in range(len(keyword_list)):
        index = main_center_classes.index(class_list[i])
        circle = plt.Circle(final_centers[index][-1], radius[i], color=colors[i], fill=True)
        del final_centers[index][-1]
        plt.gca().add_artist(circle)
        plt.text(circle.center[0], circle.center[1], keyword_list[i],
                 verticalalignment='center', horizontalalignment='center', fontname='SimHei', fontsize=10, color="white")
    plt.xlim(-5, 25)
    plt.ylim(-5, 25)
    plt.axis('off')
    # plt.show()
    plt.savefig("../front/static/new_word_cloud.png")
    print("生成词云图成功")

if __name__=='__main__':
    keywords = ['设计', '功能', '防滑性', '安全性','价格']
    importances = [1, 2, 3, 4, 5]
    classes = [1, 2, 2, 2 ,3]
    # generate_word_cloud(keywords, importances, classes)
    generate_word_cloud(keywords, importances)
